import pygame, sys, threading
from pygame import font as pf
from pygame.locals import *
from vec2d import *
from math import sqrt, e, log
from random import uniform, normalvariate
from threading import Thread
#import psyco

#single node
class node:
	def __init__(self, x, text="tag"):
		self.x = x #position
		self.oldx = x #old position
		self.a = 0 #force accumulators
				
		# Tag information
		self.text = text
		self.weight = int(round(uniform(0,100)))
		
		self.marked= False #for bfs
		self.color = (int(uniform(0,150)), int(uniform(0,150)), int(uniform(0,150)))
		self.numNodes=0 #num of nodes this is connected to. filled in countNodes()
		
#single spring
class spring:
	def __init__(self, n1, n2, k = -40, restDist = 200):
		self.n1 = n1
		self.n2 = n2
		self.k = k
		self.rest = restDist
		
#main class
class graphMain(Thread):
	def __init__(self):
		Thread.__init__(self)
		self.running = True
		
		self.nodeGetter = None
		self.nodeLock = threading._allocate_lock()
		self.updateLock = threading._allocate_lock()
		
		self.topN = []
		
		self.w = 1200
		self.h = 700
		pygame.init()
		self.screen = pygame.display.set_mode((self.w, self.h))
		self.screen.fill((0,0,0))
		
		#PHYSICS
		self.dt = 0.01 #time step
		self.friction = 0.05
		
		#lists of nodes and springs
		self.nodes = {}
		self.springs = []		
		self.dragging = None #is user dragging a node?
		
		self.selected = None #selected node
		self.physics = True #is physics enabled?
		
		self.c = 0 #program counter used in main loop
		
		self.tinyFont = pf.Font(None,20)
		self.smallFont = pf.Font(None,30)
		self.mediumFont = pf.Font(None,40)
		self.largeFont = pf.Font(None,60)
	
	#returns the top N tags
	def getTopN(self, number):
		topN = []		
		
		with self.nodeLock:
			l = self.nodes.values()
		l.sort(lambda a, b: cmp(a.weight, b.weight), reverse=True)
		
		for i in range(0, min(number, len(l))):
			topN.append(l[i])

		return topN
	
	#integrate one single time step to get new positions 
	#from old positions based on acceleration of each node.
	def verlet(self):
		for n in self.topN:
			temp = vec2d(n.x.x, n.x.y) #store old position
			# cap movement speed
			n.a.length = min(n.a.length, 2500)
			n.x += (1.0 - self.friction)*n.x - (1.0 - self.friction)*n.oldx + n.a*self.dt*self.dt
			n.oldx = temp
		
		for n in self.topN: #reset accelerations for next iteration
			n.a = vec2d(0,0)
			
	#accumulate all the forces between all the nodes
	def accumulate_force(self):
		
		#REPELL NODES CLOSE TO OTHER NODES
		#proportional to their separation in graph
		#Nodes are close => big repulsion
		#Nodes are far => no repulsion
		for n1 in self.topN:
			for n2 in self.topN:
				dst =n1.x.get_distance(n2.x) #distance on actual layout
				if dst > 0.0005:
					dp = n2.x - n1.x #get vector from n1->n2
					x0 = 0
					x1 = 300
					x = dst
					f1 = 100000/(0.001*x*x*x)
					f0 = 100000/(0.00001*x)*min(3000, n2.weight*n1.weight)		
					dp.length = f0 + (f1 - f0)/(x1 - x0)*(x-x0) #set its length to strength of interaction
					n2.a += dp #add that vector to both acceleration of n1 and n2
					n1.a += -dp	
			
			# slightly drag to center of the screen. nodes with higher weight are dragged more
			center = vec2d(self.w/2,self.h/2)
			
			dst = center.get_distance(n1.x)
			if dst > 0.0005:
					dc = center - n1.x #get vector from n1->center
					dc.length = 0.01*dst*dst*n1.weight*n1.weight*n1.weight*n1.weight
					#set its length to strength of interaction
					n1.a += dc #add that vector to both acceleration of n1 and n2
			
		#SPRING STUFF
		for s in self.springs:
			dp = s.n2.x - s.n1.x #get vector pointing from n1->n2
			
			if dp.length != 0:
				dx = dp.length - s.rest #get the displacement from rest length
				dp.length = s.k*dx #multiply by spring contant
				s.n2.a += dp #add the vector to each n1 and n2 accelerations
				s.n1.a += -dp
			
	#return the net movement of all nodes.
	#if there is very little movement then the simulation
	#can be stopped and result returned
	def netMovement(self):
		a=0.0
		for n in self.topN:
			a+= (n.x - n.oldx).length
		return a
		
	# removes nodes which have lost too much weight
	def removeDead(self):
		toRem = []
		deleted = False
		with self.nodeLock:
			for n in self.nodes.values():
				if n.weight < 0:
					for s in self.springs:
						if s.n1 == n or s.n2 == n:
							toRem.append(s)				
					self.nodes.pop(n.text)
					deleted = True
			if deleted:
				for t in toRem:
					self.springs.remove(t)
				self.doBFS()
				self.doCount()
				
	#handle all user input and interactivity
	def handle_input(self):
		for event in pygame.event.get():
			if event.type == QUIT:
				self.quit()
			elif event.type == KEYDOWN:
				if event.key == K_ESCAPE:
					self.quit()
				elif event.key == K_r:
					self.initNodes()
					self.doBFS()
					self.doCount()
				elif event.key == K_d:
					#delete closest node
					d = self.findclosest(pygame.mouse.get_pos())
					toRem = []
					for s in self.springs:
						if s.n1 == d or s.n2 == d:
							toRem.append(s)
					for t in toRem:
						self.springs.remove(t)
					self.nodes.pop(d.text)
					self.doBFS()
					self.doCount()
				elif event.key == K_p:
					self.physics = not self.physics
				elif event.key == K_n:
					with self.nodeLock:
						for z in self.nodes.values():
							z.x = vec2d(uniform(self.w/2-10,self.w/2+10), uniform(self.h/2-10,self.h/2+10))
							z.oldx = z.x
					
			elif event.type == MOUSEBUTTONUP:
				if event.button == 1:
					self.dragging = None
				else:
					now = vec2d(event.pos)
					then = vec2d(self.selected)
					if now.get_distance(then) < 10:
						#just make a new node here (at now)
						#self.nodes["tag%s%s",now] = node(now)
						#self.doBFS()
						#self.doCount()
						pass
					else:
						#make new line btw last node and this node
						#nowNode=self.findclosest(now)
						#thenNode=self.findclosest(then)
						#self.springs.append(spring(nowNode, thenNode))
						#self.doBFS()
						#self.doCount()
						pass
											
					self.selected = None
					
			elif event.type == MOUSEBUTTONDOWN:
				if event.button == 1:
					self.dragging = self.findclosest(event.pos)
				else:
					self.selected = event.pos
					
	#find the closest node to position p. p is a vec2d
	def findclosest(self, p):
		ld = self.w + self.h
		li = None
		v = vec2d(p)
		for n in self.topN:
			d = v.get_distance(n.x) 
			if d < ld:
				ld = d
				li = n
		return li
		
	#draw all springs and nodes
	def draw(self):
		
		self.screen.fill((255,255,255))
		
		# edges
		for s in self.springs:
			pygame.draw.aaline(self.screen, (200,200,200), s.n1.x.inttup(), s.n2.x.inttup(), 1)

		#with self.nodeLock:
		#	topN = self.topN
			
		for n in self.topN:
			
			# determine font size by node weight
			if (n.weight < 25):
				font = self.tinyFont
			elif(n.weight <50):
				font = self.smallFont
			elif(n.weight <75):
				font = self.mediumFont
			else:
				font = self.largeFont
			
			#pygame.draw.circle(self.screen, n.color, n.x.inttup(), 10)
			
			# text="(%s|%s)" % (x,y)
			text = "%s" % n.text
			fontFace = font.render(text, False, (0,0,0))
			
			tW, tH = font.size(text)
			
			rW = tW+10
			rH = tH+5
			
			rA = n.x.inttup()[0]-rW/2
			rB = n.x.inttup()[1]-rH/2
			tA = n.x.inttup()[0]-tW/2
			tB = n.x.inttup()[1]-tH/2
			
			pygame.draw.rect(self.screen, (255,255,255), (rA,rB,rW,rH))
			pygame.draw.rect(self.screen, (0,0,0), (rA,rB,rW,rH), 1)
			
			self.screen.blit(fontFace,(tA,tB))
		
		
		pygame.display.flip()
	
	#true if there is no spring yet between the k'th and l'th node
	def dnc(self, k, l):
		i=self.nodes.values()[k]
		j=self.nodes.values()[l]
		for s in self.springs:
			if (s.n1==i and s.n2==j) or (s.n1==j and s.n2==i):
				return False
		return True
	
	#initialize all the nodes and springs
	def initNodes(self):
		self.nodes={}
		self.springs=[]
							
		self.c = 0
		
	#do a BFS search to create the dictionary containing the length of the
	#shortest path between any two nodes in the graph, along springs
	def doBFS(self):
		
		#construct the initial dict
		self.paths={}
		for n in self.nodes.values():
			for n2 in self.nodes.values():
				if n == n2:
					self.paths[(n,n2)] = 0
				else:
					self.paths[(n,n2)] = 10 #10 if they arent connected
			
		#now run BFS on each node to find the right values and fill them in
		for n in self.nodes.values():
			lst=[n]
			d=1
			for n2 in self.nodes.values(): #reset nodes
				n2.marked = False
				
			while len(lst) > 0:
				nn = lst[0]
				del lst[0]
				nn.marked = True
				
				for s in self.springs:
					n2 = None
					if s.n1 == nn:
						n2 = s.n2
					elif s.n2 == nn:
						n2 = s.n1
						
					if n2 != None and n2.marked == False:
						#nn is part of this spring
						#we found n2 in depth d from n
						self.paths[(n,n2)] = min(d, self.paths[(n,n2)])
						self.paths[(n2,n)]= min(d, self.paths[(n2,n)])
						n2.color = nn.color
						
						#append n2 to the lst
						lst.append(n2)
						
				d+=1 #increment depth
	
	#count number of connections for each node
	#return the highest number of connections of any node
	def doCount(self):
		maxN = 0
		
		for n in self.topN:
			k=0
			for s in self.springs:
				if s.n1 == n or s.n2 == n:
					k+=1
			n.numNodes=k/2 #k/2 because graph is undirected. nodes are doublecounted
			if k/2 > maxN:
				maxN = k/2
		
		#now set the spring constants appropriately
		for s in self.springs:
			n=max(s.n1.numNodes, s.n2.numNodes)
			#s.rest=50+n*10
			
		return maxN
	
	def updateNodes(self, newNodes=[]):
		with self.updateLock:
			for new in newNodes:
				try: 
					with self.nodeLock:
						self.nodes[new[0]].weight += new[1]
				except:
					# print "!!! creating new node for tag: %s, %s" % (new[0], new[1])
					# create new node
					z=node(vec2d(uniform(self.w/2-100,self.w/2+100), uniform(self.h/2-100,self.h/2+100)))
					z.text = new[0]
					z.weight = new[1]
					with self.nodeLock:
						self.nodes[z.text] = z
			
	#main method. do the simulation
	def run(self):	
		
		#initialize the particle system
		with self.nodeLock:
			self.initNodes()
			self.doBFS()
			self.doCount()
		
		try:
			self.nodeGetter.start()
		except:
			print "No nodeGetter could be started."		
		
		while self.running:
			
			self.c +=1			
			
			# update topN every 100 frames
			if self.c % 100 == 0:
				# recalculate topN
				topN = self.getTopN(20)
				
				with self.nodeLock:
					self.topN = topN
			

			#if physics is enabled then do simulation
			if self.physics: 
				with self.nodeLock:
					self.accumulate_force()
					self.verlet()
			
			self.handle_input() #handle all user input
			
			#handle node dragging
			if not self.dragging == None:
				pos = pygame.mouse.get_pos()
				self.dragging.x = vec2d(pos)
				self.dragging.oldx = vec2d(pos)
						
			#draw everything
			self.draw()
			
	def quit(self):
		self.running = False
		self.nodeGetter.stop()
		sys.exit()